import pandas as pd
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression
import numpy as np

df = pd.read_csv('static/data/info_pre.csv')
df_word = pd.read_csv('static/data/word.csv')


def album_count_top5():
    df_album_top5 = df.drop_duplicates(subset='title', keep='first', inplace=False, ignore_index=False)
    df_album_top5 = df_album_top5.sort_values(by='album_count', ascending=False).head(5)
    name = df_album_top5['title'].values.tolist()
    value = df_album_top5['album_count'].values.tolist()
    return {
        'name': name,
        'value': value
    }


def category_ratio():
    df_category_ratio = df[['category', 'title']].groupby('category').count().reset_index().sort_values(by='title',
                                                                                                        ascending=False)
    data = [{'value': int(i[1]), 'name': i[0]} for i in df_category_ratio.values]
    return {
        'data': data
    }


def scatter_data():
    scatter_list = df[['album_count', 'play_count']].values.tolist()
    return {
        'data': scatter_list
    }


def treemap_data():
    rank_type_list = set(df['rank_type'].values)
    data = [
        {
            'name': i,
            'value': len(df[df['rank_type'] == i]),
            'children': [
                {
                    'name': j[0],
                    'value': len(j[1]),
                    'children': [
                        {
                            'name': f'{k[1]}-{k[0]}',
                            'value': 1
                        }
                        for k in j[1].sort_values(by='rank', ascending=False).values.tolist()
                    ]
                } for j in df[df['rank_type'] == i].groupby('category')[['title', 'rank']]
            ]
        } for i in rank_type_list
    ]
    return {
        'data': data
    }


def word_data():
    data = [
        {
            'name': i[0],
            'value': int(i[1])
        } for i in df_word.head(50).values
    ]
    return {
        'data': data
    }
